package com.wang.tool.util;


import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLDecoder;
import java.util.Enumeration;
import java.util.function.Consumer;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.stream.Stream;

public class ClassScannerUtil {


    public static void scan(String packageName, Consumer<Class<?>> consumer){

        ClassLoader classLoader =  Thread.currentThread().getContextClassLoader();
        Enumeration<URL> resources = null;
        try {
            resources = classLoader.getResources(packageName.replace(".", "/"));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        while (resources.hasMoreElements()) {
            URL resource = resources.nextElement();
            if ("file".equals(resource.getProtocol())) {
                scanClassesInDirectory(resource.getFile(), packageName, consumer);
            } else if ("jar".equals(resource.getProtocol())) {
                JarFile jar = null;
                try {
                    jar = new JarFile(URLDecoder.decode(resource.getFile(), "UTF-8"));
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
                scanClassesInJar(jar, packageName, consumer);
            }
        }


    }

    private static void scanClassesInDirectory(String directoryPath, String packageName, Consumer<Class<?>> consumer) {
        File directory = new File(directoryPath);
        Stream.of(directory.listFiles())
              .filter(file -> file.isFile() && file.getName().endsWith(".class"))
              .map(file -> file.getName().substring(0, file.getName().length() - 6))
              .forEach(className -> {
                  try {
                      consumer.accept(Class.forName(packageName + "." + className));
                  } catch (ClassNotFoundException e) {
                      throw new RuntimeException(e);
                  }
              });
    }

    private static void scanClassesInJar(JarFile jar, String packageName, Consumer<Class<?>> consumer) {
        Enumeration<JarEntry> entries = jar.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (!entry.isDirectory() && entry.getName().startsWith(packageName.replace(".", "/")) && entry.getName().endsWith(".class")) {
                String className = entry.getName().replaceAll("/", ".");
                className = className.substring(0, className.length() - 6);
                try {
                    consumer.accept(Class.forName(className));
                } catch (ClassNotFoundException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

}